import os
from datetime import datetime

from loguru import logger
from torchtext.datasets import AG_NEWS
from torchtext.vocab import build_vocab_from_iterator
from tqdm import tqdm


def load_dataset(dataset_file_path, split):
    """
    :description： 根据指定参数，加载数据集
    :param dataset_file_path: 数据集存储路径
    :param split: 加载数据集类别，split=train，加载训练集；split=test，加载测试集
    :return:
    """
    logger.debug(f'start loading {split} dataset')
    data = AG_NEWS(root=dataset_file_path, split=split)
    logger.debug(f'loading {split} dataset successfully')
    return data


def build_vocab(data_file_path, split, tokenizer):
    """
    :description: 使用指定数据集全量数据，构建词汇表 vocab
    :param data_file_path: 数据集存储路径
    :param split: 加载数据集类别，split=train，加载训练集；split=test，加载测试集
    :param tokenizer: 文本转token工具，可理解为分词器
    :return:
    """
    data = load_dataset(data_file_path, split)
    logger.debug('start building vocab')
    tokens = []
    for _, text in tqdm(list(data), ncols=100):
        tokens.append(tokenizer(text))
    vocab = build_vocab_from_iterator(iter(tokens), specials=["<unk>"])
    vocab.set_default_index(vocab["<unk>"])
    logger.debug('building vocab successfully')
    return vocab


def log_init(file_name, log_dir='../logs'):
    """
    :description: 日志对象初始化配置函数
    :param file_name: 日志文件名
    :param log_dir: 日志文件保存路径
    :return:
        loguru.logger
    """
    retention_ = "1 days"
    encoding_ = "utf-8"

    # 设置日志保存目录
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    log_path = os.path.join(log_dir, file_name + '_' + str(datetime.now())[:10] + '.txt')

    # 配置logger对象
    logger.add(sink=log_path, rotation=retention_, encoding=encoding_)
